package org.kiwi.data.adapter.druid;

import com.alibaba.druid.pool.DruidConnectionHolder;
import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.DruidPooledConnection;
import org.apache.commons.lang3.StringUtils;
import org.kiwi.data.TableShardingStrategyHolder;
import org.kiwi.data.vo.TableShardingParameter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cglib.proxy.Enhancer;
import org.springframework.cglib.proxy.MethodInterceptor;
import org.springframework.cglib.proxy.MethodProxy;
import org.springframework.core.LocalVariableTableParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;

import java.lang.reflect.Method;
import java.sql.SQLException;

import static org.kiwi.data.util.ShardingUtil.convertSql;

/**
 * Created by jack08.liu on 2016/12/22.
 */
public class DruidDataSourceAdapter extends DruidDataSource {
    private static final Logger LOGGER = LoggerFactory.getLogger(DruidDataSourceAdapter.class);

    private static final String MODE_NORMAL = "normal";
    private static final String MODE_ADAPTER = "adapter";
    private static final String MODE_PROXY = "proxy";

    private static final String DEFAULT_MODE = MODE_NORMAL;

    private String mode = DEFAULT_MODE;

    public void setMode(String mode) {
        this.mode = mode;
    }

    @Override
    public DruidPooledConnection getConnection() throws SQLException {
        if (StringUtils.equalsIgnoreCase(MODE_NORMAL, this.mode)) {
            return super.getConnection();
        } else if (StringUtils.equalsIgnoreCase(MODE_ADAPTER, this.mode)) {
            return getConnection(this.maxWait);// TODO: 2016/12/24 return DruidPooledConnectionAdapter
        } else if (StringUtils.equalsIgnoreCase(MODE_PROXY, this.mode)) {
            return createDruidPooledConnectionProxy(super.getConnection());
        }
        throw new IllegalArgumentException("invalid connection mode");
    }

    private DruidPooledConnection createDruidPooledConnectionProxy(DruidPooledConnection rawConnection) {
        Enhancer enhancer = new Enhancer();
        enhancer.setSuperclass(DruidPooledConnection.class);
        enhancer.setCallback(new DruidPooledConnectionInterceptor(rawConnection));
        return (DruidPooledConnection) enhancer.create(new Class[]{DruidConnectionHolder.class},
                new Object[]{rawConnection.getConnectionHolder()});
    }

    private static class DruidPooledConnectionInterceptor implements MethodInterceptor {

        public static final String SPECIFY_PARAMETER_NAME = "sql";
        public static final String CALL = "Call";
        public static final String STATEMENT = "Statement";

        private DruidPooledConnection rawConnection;

        public DruidPooledConnectionInterceptor(DruidPooledConnection rawConnection) {
            this.rawConnection = rawConnection;
        }

        @Override
        public Object intercept(Object o, Method method, Object[] args, MethodProxy methodProxy) throws Throwable {
            try {
                TableShardingParameter tableShardingParameter = TableShardingStrategyHolder.getTableShardingParameter();

                if (tableShardingParameter != null && StringUtils.containsAny(method.getName(), STATEMENT, CALL)) {
                    ParameterNameDiscoverer parameterNameDiscoverer = new LocalVariableTableParameterNameDiscoverer();
                    String[] parameterNames = parameterNameDiscoverer.getParameterNames(method);
                    if (parameterNames != null && parameterNames.length > 0) {
                        for (int i = 0; i < parameterNames.length; i++) {
                            if (StringUtils.equalsIgnoreCase(SPECIFY_PARAMETER_NAME, parameterNames[i])) {
                                args[i] = convertSql(args[i], tableShardingParameter);
                                if (LOGGER.isDebugEnabled()) {
                                    LOGGER.debug("after table sharding , sql=[{}]", args[i]);
                                }
                                break;
                            }
                        }
                    }
                }

                return method.invoke(this.rawConnection, args);

            } catch (Throwable e) {
                throw e;
            }
        }
    }
}
